# coding=utf-8
import zipfile
import os
import cv2
import numpy as np

from io import BytesIO


class ZipFileReader:
    def __init__(self, bytes):
        self.f = BytesIO()
        self.f.write(bytes)
        self.z = zipfile.ZipFile(self.f, 'r')
        self.file_paths = [info.filename for info in self.z.infolist() if not info.is_dir()]
        self.num = len(self.file_paths)

    def __iter__(self):
        self.count = 0
        return self

    def __next__(self):
        if self.count == self.num:
            self.z.close()
            raise StopIteration
        with self.z.open(self.file_paths[self.count], 'r') as f:
            img_bytes = f.read()
            img_array = cv2.imdecode(np.frombuffer(img_bytes, np.uint8), cv2.IMREAD_COLOR)

        self.count += 1
        return img_array

    def __len__(self):
        return self.num


if __name__ == '__main__':
    path = '../test_zip/val.zip'
    with open(path, 'rb') as f:
        file = f.read()
    zf_iter = ZipFileReader(file)
    for img in zf_iter:
        print(img)
